from datetime import datetime, timedelta
from typing import List, Optional
from pytz import timezone

# from hs_udata import set_token, fut_quote_minute, stock_quote_minutes, fut_list
from akshare import stock_zh_a_hist, stock_individual_info_em
from pandas import DataFrame
import pandas as pd

from vnpy.trader.setting import SETTINGS
from vnpy.trader.constant import Exchange, Interval
from vnpy.trader.object import BarData, HistoryRequest, FutBasicData, FutBasicRequest
from vnpy.trader.datafeed import BaseDatafeed
from time import sleep


EXCHANGE_VT2UDATA = {
    Exchange.CFFEX: "CFE",
    Exchange.SHFE: "SHF",
    Exchange.DCE: "DCE",
    Exchange.CZCE: "CZC",
    Exchange.INE: "INE",
    Exchange.SSE: "SH",
    Exchange.SZSE: "SZ"
}

INTERVAL_VT2RQ = {
    Interval.MINUTE: "1m",
    Interval.HOUR: "60m",
    Interval.DAILY: "1d",
}

INTERVAL_ADJUSTMENT_MAP = {
    Interval.MINUTE: timedelta(minutes=1),
    Interval.HOUR: timedelta(hours=1),
    Interval.DAILY: timedelta()         # no need to adjust for daily bar
}


CHINA_TZ = timezone("Asia/Shanghai")


def convert_symbol(symbol: str, exchange: Exchange) -> str:
    """将交易所代码转换为Akshare代码"""
    exchange_str = EXCHANGE_VT2UDATA.get(exchange, "")
    return f"{symbol.upper()}.{exchange_str}"


def convert_exchange(exchange: Exchange) -> str:
    """将交易所代码转换为Akshare代码"""
    return EXCHANGE_VT2UDATA.get(exchange, "")



class AKShareDatafeed(BaseDatafeed):
    """AKShare数据服务接口"""
    def __init__(self):
        """"""
        # self.token: str = SETTINGS["datafeed.password"]

    def query_bar_history(self, req: HistoryRequest) -> Optional[DataFrame]:
        # 只支持分钟线
        # 历史数据：2017-08-01至今
        # if req.interval != Interval.MINUTE:
        #     return None
        interval = req.interval

        data: DataFrame = pd.DataFrame()

        end = req.end
        if req.interval == Interval.DAILY:
            once_day = 365
        elif req.interval == Interval.MINUTE:
            once_day = 60

        # 注：起始日期与结束日期间隔应小于60天，单次调用最大返回条数为10000条；更多数据可对股票代码与日期进行循环
        temp_end = req.start + timedelta(days=once_day)

        if end > temp_end:
            req.end = temp_end

        while True:
            # 期货
            if req.exchange in {
                Exchange.CFFEX,
                Exchange.SHFE,
                Exchange.CZCE,
                Exchange.DCE,
                Exchange.INE
            }:
                temp_data = self.query_futures_bar_history(req)
                if not temp_data:
                    return data
                # data.extend(temp_data)
                data.append(temp_data)

                # if temp_data[-1].datetime.date() >= end or len(temp_data) != 10000:
                #     break
                if temp_data.loc[-1, 'datetime'] >= end:
                    break
                req.start = temp_data[-1].datetime
            # 股票
            elif req.exchange in {
                Exchange.SSE,
                Exchange.SZSE
            }:
                temp_data = self.query_equity_bar_history(req)
                print(len(data))
                if len(temp_data) == 0:
                    return data
                data = pd.concat([temp_data, data])
                req.start = req.end
                req.end = req.start + timedelta(days=once_day)
                temp_end = temp_data['datetime'].iloc[-1]
                if temp_end >= end or req.end >= end:
                    break
            # 其他
            else:
                return None
        return data


    def query_futures_bar_history(self, req: HistoryRequest) -> Optional[List[BarData]]:
        """查询期货分钟K线数据"""
        symbol = req.symbol
        exchange = req.exchange
        interval = req.interval
        start = req.start
        end = req.end

        udata_symbol = convert_symbol(symbol, exchange)
        adjustment = timedelta(minutes=1)

        df: DataFrame = fut_quote_minute(
            en_prod_code=udata_symbol,
            begin_date=start.strftime("%Y-%m-%d"),
            end_date=end.strftime("%Y-%m-%d")
        )

        data: List[BarData] = []

        if len(df):
            for _, row in df.iterrows():
                timestr = f"{row.date} {str(row.time).rjust(4, '0')}"
                dt = datetime.strptime(timestr, "%Y-%m-%d %H%M") - adjustment
                dt = CHINA_TZ.localize(dt)

                bar = BarData(
                    symbol=symbol,
                    exchange=exchange,
                    interval=interval,
                    datetime=dt,
                    open_price=row.open,
                    high_price=row.high,
                    low_price=row.low,
                    close_price=row.close,
                    volume=row.turnover_volume,
                    turnover=row.turnover_value,
                    open_interest=row.amount,
                    gateway_name="UDATA"
                )

                data.append(bar)

        return data


    def query_equity_bar_history(self, req: HistoryRequest) -> Optional[DataFrame]:
        """查询股票分钟K线数据"""
        symbol = convert_symbol(req.symbol, req.exchange)
        exchange = req.exchange
        interval = req.interval
        start = req.start
        end = req.end
        # if (interval == Interval.HOUR):
        #     interval = Interval.MINUTE
        # udata_symbol = convert_symbol(symbol, exchange)
        # adjustment = timedelta(minutes=1)

        df: DataFrame = stock_zh_a_hist(
            symbol=req.symbol,
            period="daily",
            start_date=start.strftime("%Y%m%d"),
            end_date=end.strftime("%Y%m%d")
        )

        if len(df):
            df['symbol'] = req.symbol
            df['exchange'] = req.exchange
            df["open"] = pd.to_numeric(df['开盘'])
            df["close"] = pd.to_numeric(df['收盘'])
            df["low"] = pd.to_numeric(df['最低'])
            df["high"] = pd.to_numeric(df['最高'])
            df["volume"] = pd.to_numeric(df['成交量'])
            df["turnover"] = pd.to_numeric(df['成交额'])
            df['datetime'] = pd.to_datetime(df['日期'], format="%Y-%m-%d", utc=False)
            df['gateway_name'] = "AKShare"
            df['open_interest'] = 0
            df['interval'] = interval
            # timestr = f"{row.date} {str(row.time).rjust(4, '0')}"
            # dt = datetime.strptime(timestr, "%Y-%m-%d %H%M") - adjustment
            # dt = CHINA_TZ.localize(dt)
            print(df.head().to_string())

        return df

    def query_fut_basic(self, req: FutBasicRequest) -> Optional[List[FutBasicData]]:

        """获取期货合约列表数据"""
        exchange = req.exchange
        fut_type = req.fut_type

        df: DataFrame = fut_list()

        data: List[FutBasicData] = []

        if len(df):
            for _, row in df.iterrows():
                fut = FutBasicData(
                    symbol=row.secu_code,
                    exchange=Exchange.SSE,
                    name=row.secu_abbr,
                    gateway_name="UDATA"
                )

                data.append(fut)

        return data
